#python -m lilab.OpenLabCluster_train.a6c_clippredpkl_merge CLIPREDPKL_FILE1 CLIPREDPKL_FILE2 ... CLIPREDPKL_FILEN
#%%
import pickle
import numpy as np
import os
import os.path as osp
import argparse
import pandas as pd


clippredfiles = [
    '/DATA/zhongzhenchao/24maternalBehv/usv-data/redo3dbehv/out_semiseq2seq_norm/lstm_offline.clippredpkl',
    '/mnt/liying.cibr.ac.cn_Data_Temp_ZZC2/0425batch-chr2lightCtrl/all/out_semiseq2seq_norm/lstm_offline.clippredpkl'
]


def main(clippredfiles, outdir=None):
    assert len(clippredfiles)>=2, 'at least 2 clippredfiles are needed'
    assert all(osp.exists(clippredfile) for clippredfile in clippredfiles), 'clippredfile not exist'
    if outdir is None:
        outdir = osp.join(osp.dirname(clippredfiles[0]), 'clip_merges')
    os.makedirs(outdir, exist_ok=True)
    outfile = osp.join(outdir, osp.basename(clippredfiles[0]))

    # %%
    keys_to_cp = ['ncluster', 'cluster_names', 'nK_mutual', 'ntwin', 'clipNames', 'nK_mirror_half', 'cluster_names_mutualmerge']
    clippreddata_l = [pickle.load(open(clippredfile, 'rb')) for clippredfile in clippredfiles]
    assert len(set([d['ncluster'] for d in clippreddata_l])) == 1, 'ncluster is not the same'
    assert len(set([tuple(d['cluster_names']) for d in clippreddata_l])) == 1, 'cluster_names is not the same'

    outdict = {key: clippreddata_l[0][key] for key in keys_to_cp}
    outdict['df_clipNames'] = pd.concat([d['df_clipNames'] for d in clippreddata_l]).reset_index(drop=True)
    has_duplicated = outdict['df_clipNames'][['vnake', 'startFrame', 'isBlack']].duplicated().any()
    assert not has_duplicated, 'duplicated clipNames'

    for key in ['embedding_d2', 'cluster_labels', 'cluster_pvalues', 'clipNames']:
        outdict[key] = np.concatenate([d[key] for d in clippreddata_l], axis=0)

    try:
        outdict['embedding'] = np.concatenate([d['embedding'] for d in clippreddata_l], axis=0)
    except:
        print('embedding not found or not the same shape')

    pickle.dump(outdict, open(outfile, 'wb'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Merge clippredpkl files')
    parser.add_argument('clippredfiles', nargs='+', help='clippredpkl files to merge')
    parser.add_argument('--outdir', default=None, help='output directory')
    args = parser.parse_args()
    main(args.clippredfiles, args.outdir)
